package com.lw.smartquiz.manager;

import com.lw.smartquiz.constant.AIConstant;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.*;
import io.reactivex.Flowable;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.stereotype.Component;

import javax.annotation.Resource;
import java.util.ArrayList;
import java.util.List;

/**
 * 封装AI使用方法，方便请求智谱AI
 */
@Slf4j
@Component
public class AIManager {
    @Resource
    private ClientV4 clientV4;

    @Value("${ai.max-tokens}")
    private int MAX_TOKENS;

    /**
     * 流式传输请求，回答较随机
     *
     * @param systemMessage 系统消息
     * @param userMassage   用户消息
     * @return AI回答结果
     */
    public Flowable<ModelData> streamUnstableRequest(String systemMessage, String userMassage){
        return streamRequest(systemMessage, userMassage, AIConstant.UNSTABLE_TEMPERATURE);
    }

    /**
     * 流式传输请求，回答较稳定
     *
     * @param systemMessage 系统消息
     * @param userMassage   用户消息
     * @return AI回答结果
     */
    public Flowable<ModelData> streamStableRequest(String systemMessage, String userMassage){
        return streamRequest(systemMessage, userMassage, AIConstant.STABLE_TEMPERATURE);
    }

    /**
     * 非流式传输请求，回答较随机
     * @param systemMessage 系统消息
     * @param userMassage 用户消息
     * @return AI回答结果
     */
    public String syncUnstableRequest(String systemMessage, String userMassage){
        return syncRequest(systemMessage, userMassage, AIConstant.UNSTABLE_TEMPERATURE);
    }

    /**
     * 非流式传输请求，回答较稳定
     * @param systemMessage 系统消息
     * @param userMassage 用户消息
     * @return AI回答结果
     */
    public String syncStableRequest(String systemMessage, String userMassage){
        return syncRequest(systemMessage, userMassage, AIConstant.STABLE_TEMPERATURE);
    }

    /**
     * 流式传输请求
     *
     * @param systemMessage 系统消息
     * @param userMassage   用户消息
     * @return AI回答结果
     */
    public Flowable<ModelData> streamRequest(String systemMessage, String userMassage, Float temperature){
        ModelApiResponse modelApiResponse = doRequest(systemMessage, userMassage, Boolean.TRUE, temperature);
        return modelApiResponse.getFlowable();
    }

    /**
     * 同步请求
     * @param systemMessage 系统消息
     * @param userMessage   用户消息
     * @param temperature   随机性，值越大，回答越不稳定
     * @return AI回答结果
     */
    public String syncRequest(String systemMessage, String userMessage, Float temperature){
        //请求
        ModelApiResponse modelApiResponse = this.doRequest(systemMessage, userMessage, Boolean.FALSE, temperature);
        //返回响应内容
        String result = modelApiResponse.getData().getChoices().get(0).getMessage().getContent().toString();
        log.debug("AI回答内容:\n{}", result);
        return result;
    }

    /**
     * 传递系统消息和用户消息
     *
     * @param systemMessage 系统消息
     * @param userMessage   用户消息
     * @param stream        是否流式传输
     * @param temperature   随机性，值越大，回答越不稳定
     * @return AI回答结果
     */
    public ModelApiResponse doRequest(String systemMessage, String userMessage, Boolean stream, Float temperature){
        //构建信息
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage systemChatMessage = new ChatMessage(ChatMessageRole.SYSTEM.value(), systemMessage);
        ChatMessage userChatMessage = new ChatMessage(ChatMessageRole.USER.value(), userMessage);
        messages.add(systemChatMessage);
        messages.add(userChatMessage);
        log.debug("系统提示词：\n{}", systemMessage);
        log.debug("用户提示词：\n{}", userMessage);
        return doRequest(messages, stream, temperature);
    }

    /**
     * 通用请求
     *
     * @param messages    消息列表
     * @param stream      是否流式传输
     * @param temperature 随机性，值越大，回答越不稳定
     * @return AI回答结果
     */
    public ModelApiResponse doRequest(List<ChatMessage> messages, Boolean stream, Float temperature){
        //构建请求
        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM4) //使用的模型
                .stream(stream) //流式传输
                .invokeMethod(Constants.invokeMethod) //同步调用
                .temperature(temperature) //随机性
                .maxTokens(MAX_TOKENS) //最大token
                .messages(messages)
                .build();
        //调用接口，返回
        return clientV4.invokeModelApi(chatCompletionRequest);
    }
}
